2.3 S4(Structured State Spaces) 대각화와 컨볼루션 기반 학습 가속

2.3 S4(Structured State Spaces) 대각화와 컨볼루션 기반 학습 가속

상태 공간 모델(SSM)이 딥러닝의 주류 아키텍처로 부상하기 위해 넘어야 했던 가장 거대한 장벽은 이론적 우아함과 계산적 효율성 사이의 괴리였다. 앞선 절에서 논의한 HiPPO(High-order Polynomial Projection Operator) 이론은 순환 신경망이 장기 의존성(Long-Range Dependency)을 학습할 수 있는 수학적 토대를 마련해주었으나, 이를 실제 하드웨어에서 구현하는 것은 전혀 다른 차원의 문제였다. HiPPO 행렬은 그 구조적 특성상 비정규(Non-normal) 행렬이며, 이는 일반적인 대각화 기법을 거부한다. 본 절에서는 이러한 난제를 해결하고 SSM을 트랜스포머보다 빠른 학습 및 추론 모델로 탈바꿈시킨 핵심 기술, 즉 **S4(Structured State Space Sequence Model)**의 대각화 기법과 컨볼루션 기반 가속 알고리즘에 대해 심층적으로 기술한다.

Gu et al.(2021)에 의해 제안된 S4 모델은 연속 시간 시스템의 이점과 이산 시간 연산의 효율성을 결합하기 위해 선형대수학의 정교한 정리들을 집대성했다. 특히 NPLR(Normal Plus Low-Rank) 분해, 우드버리 행렬 항등식(Woodbury Matrix Identity), 그리고 **코시 커널(Cauchy Kernel)**로 이어지는 논리적 전개는 O(N^3)의 연산 비용이 소요되는 행렬 역변환 과정을 O(N \log N) 또는 O(L \log L) 수준으로 낮추는 획기적인 전환점이 되었다. 이 기술적 혁신이 없었다면 오늘날의 Mamba 아키텍처는 존재할 수 없었을 것이다. 따라서 본 절에서는 S4의 수학적 엔진을 구성하는 각 요소를 해부하고, 이것이 어떻게 긴 문맥(Long Context) 처리를 위한 선형 시간 복잡도를 달성했는지 규명한다.

1. 계산 복잡도의 딜레마와 대각화의 필요성

시퀀스 모델링에서 상태 공간 모델은 다음의 선형 미분 방정식으로 정의된다.
\begin{aligned} x'(t) &= Ax(t) + Bu(t) \\ y(t) &= Cx(t) + Du(t) \end{aligned}
여기서 상태 벡터 x(t) \in \mathbb{R}^N은 입력 u(t)의 과거 정보를 압축하여 저장하는 메모리 역할을 수행한다. 이 시스템을 컴퓨터에서 학습시키기 위해서는 이산화(Discretization) 과정을 거쳐야 하며, 이는 다음과 같은 점화식 형태로 나타난다.
x_k = \bar{A}x_{k-1} + \bar{B}u_k

y_k = \bar{C}x_k + \bar{D}u_k

여기서 \bar{A}, \bar{B}는 연속 시간 파라미터 A, B와 단계 크기(step size) \Delta에 의해 결정되는 이산화된 행렬들이다.

문제는 상태 차원 N과 시퀀스 길이 L이 커질 때 발생한다. 전통적인 RNN 방식(Recurrent view)으로 데이터를 처리할 경우, 시퀀스 길이가 길어질수록 그래디언트 소실 문제가 발생할 뿐만 아니라, GPU와 같은 병렬 처리 하드웨어의 이점을 살리지 못한 채 순차적으로 연산을 수행해야 하므로 학습 속도가 매우 느리다. 반면, 이를 컨볼루션 형태(Convolutional view)로 변환하여 병렬 처리를 수행하려 해도, 컨볼루션 커널 \bar{K}를 생성하는 과정에서 \bar{A}의 거듭제곱(\bar{A}, \bar{A}^2, \dots, \bar{A}^{L-1})을 계산해야 하므로 나이브한 접근 방식은 O(N^2 L)의 연산량을 요구한다.

상태 차원 N은 모델의 표현력(Expressivity)을 결정짓는 중요한 요소다. N을 충분히 키우지 않으면 복잡한 데이터의 패턴을 포착할 수 없다. 그러나 O(N^2 L)의 복잡도는 N을 키우는 것을 계산 비용 측면에서 불가능하게 만든다. 트랜스포머 모델이 O(L^2)의 복잡도를 가져 긴 시퀀스 처리에 취약하다면, 초기 SSM은 O(N^2)의 복잡도로 인해 깊은 상태 표현을 갖는 데 취약했다.

이 딜레마를 해결하는 유일한 열쇠는 상태 행렬 A를 **대각화(Diagonalization)**하는 것이다. 만약 A가 대각 행렬(Diagonal Matrix) \Lambda와 닮음(Similar) 관계라면, 즉 A = V \Lambda V^{-1}로 분해될 수 있다면, 행렬의 거듭제곱 A^kV \Lambda^k V^{-1}로 매우 쉽게 계산될 수 있다. 대각 행렬의 거듭제곱은 각 원소를 거듭제곱하는 것과 같으므로 연산 복잡도가 O(N)으로 획기적으로 줄어든다.

그러나 앞선 절에서 다룬 HiPPO 행렬은 장기 의존성을 포착하기 위해 특수한 구조를 띠고 있으며, 이는 수학적으로 비정규(Non-normal) 행렬에 해당한다. 비정규 행렬은 고유벡터(Eigenvector)들이 서로 직교하지 않는 행렬을 말한다. HiPPO 행렬을 억지로 대각화하려 할 경우, 고유벡터 행렬 V의 조건수(Condition Number, \kappa(V) = \|V\| \|V^{-1}\|)가 천문학적으로 커져 수치적 불안정성을 초래한다. 실제로 HiPPO-LegS 행렬의 경우 N이 커짐에 따라 조건수가 발산하여, 배정밀도(Double Precision) 연산 체계에서도 유효 숫자를 모두 잃어버리는 문제가 발생한다.

S4의 저자들은 이 지점에서 타협하지 않고 새로운 돌파구를 제시했다. 그들은 A를 완벽한 대각 행렬로 만드는 대신, **“대각 행렬과 저순위(Low-Rank) 행렬의 합”**으로 표현할 수 있다면 효율적인 연산이 가능함을 간파했다. 이것이 바로 S4 알고리즘의 핵심인 NPLR(Normal Plus Low-Rank) 분해이다.

2. NPLR(Normal Plus Low-Rank) 분해와 HiPPO의 재해석

S4 아키텍처의 가장 근본적인 기여는 HiPPO 행렬을 포함한 다양한 구조적 행렬들이 NPLR 형태로 표현될 수 있음을 증명하고, 이를 활용한 고속 알고리즘을 제안한 것이다.

[정의 2.3.1] NPLR 파라미터화

어떤 행렬 A \in \mathbb{C}^{N \times N}가 다음과 같이 분해될 수 있다면, 이를 NPLR 구조를 갖는다고 한다.
A = V \Lambda V^* - P Q^*
여기서:

  • V \in \mathbb{C}^{N \times N}는 유니타리 행렬(Unitary Matrix)로, VV^* = I를 만족한다.
  • \Lambda \in \mathbb{C}^{N \times N}는 대각 행렬(Diagonal Matrix)이다.
  • P, Q \in \mathbb{C}^{N \times r}는 낮은 랭크(Low-Rank)의 행렬이다 (r \ll N). SSM에서는 주로 r=1 또는 r=2인 경우가 사용된다.

이 정의가 중요한 이유는 두 가지다. 첫째, HiPPO 이론에서 유도된 모든 주요 행렬들(LegS, LagT 등)이 이 형태를 만족한다. 둘째, 유니타리 행렬 V에 의한 공액 변환(Conjugation)은 수치적으로 매우 안정적이다.

HiPPO-LegS 행렬의 경우를 살펴보자. 이 행렬 A는 다음과 같은 형태를 띤다.
A_{nk} = \begin{cases} -(2n+1)^{1/2}(2k+1)^{1/2} & n > k \\ -(n+1) & n = k \\ 0 & n < k \end{cases}
Gu et al. (2021)은 이 행렬이 A = - \frac{1}{2} P P^* - i S (여기서 S는 대칭 행렬) 형태로 분해될 수 있음을 보였으며, 이는 곧 A가 정규 행렬(Skew-Symmetric 부분)과 저순위 섭동(Low-Rank Perturbation)의 합으로 표현됨을 의미한다. 이를 정리하면 A는 대각 행렬에 랭크-1 또는 랭크-2의 수정을 가한 형태와 유니타리 닮음(Unitary Similar) 관계에 있다.

이러한 NPLR 구조의 발견은 A 행렬을 직접 다루는 대신, 이를 구성하는 파라미터 \Lambda, P, Q, V를 학습하거나 고정하여 사용할 수 있는 길을 열어주었다. 특히 S4 알고리즘은 V를 항등 행렬에 가깝게 만들거나 미리 계산된 고정 유니타리 행렬로 처리함으로써, 실질적으로 ADPLR(Diagonal Plus Low-Rank) 행렬, 즉 대각 행렬과 저순위 행렬의 합으로 취급한다.
A \approx \Lambda - P Q^*
이제 문제는 ( \Lambda - P Q^* ) 형태의 행렬을 사용하여 어떻게 효율적인 이산화와 컨볼루션 커널 생성을 수행할 것인가로 귀결된다.

3. 쌍선형 변환(Bilinear Discretization)과 구조 보존

연속 시간 시스템을 이산 시간 시스템으로 변환하는 방법에는 오일러 방법(Euler Method), 영차 홀드(ZOH), 쌍선형 변환(Bilinear Transform, 튜스틴 방법) 등이 있다. RNN 연구에서는 주로 단순한 오일러 방법을 사용해왔으나, S4는 시스템의 안정성과 주파수 응답 특성을 보존하기 위해 쌍선형 변환을 채택한다.

단계 크기 \Delta에 대해 쌍선형 변환은 다음과 같이 정의된다.
\bar{A} = (I - \frac{\Delta}{2} A)^{-1} (I + \frac{\Delta}{2} A)
이 식은 z-변환 영역에서 s \approx \frac{2}{\Delta} \frac{1-z^{-1}}{1+z^{-1}}로 치환하는 것과 등가이다. 여기서 주목해야 할 점은 행렬 역변환 (I - \frac{\Delta}{2} A)^{-1}이 포함되어 있다는 것이다. 일반적인 행렬 A에 대해 이 역행렬을 구하는 것은 O(N^3)의 비용이 든다. 매 학습 단계마다 \Delta가 학습 가능한 파라미터로서 변할 수 있음을 고려하면, 이 역행렬 연산을 반복하는 것은 막대한 계산 부담이 된다.

하지만 A가 NPLR 구조를 가진다면 이야기는 달라진다.
I - \frac{\Delta}{2} A = I - \frac{\Delta}{2} (\Lambda - P Q^*) = (I - \frac{\Delta}{2} \Lambda) + \frac{\Delta}{2} P Q^*
위 식의 우변은 대각 행렬과 저순위 행렬의 합이다. 선형대수학의 성질에 따르면, 어떤 행렬이 DPLR(Diagonal Plus Low-Rank) 형태라면 그 역행렬 또한 DPLR 형태를 유지한다. 따라서 쌍선형 변환을 거친 이산 행렬 \bar{A} 역시 NPLR 구조를 보존하게 되며, 이는 후술할 우드버리 항등식을 통한 고속 연산의 기초가 된다.

4. SSM 생성 함수와 우드버리 항등식의 마법

S4 모델의 궁극적인 목표는 전체 시퀀스에 적용될 컨볼루션 커널 \bar{K}를 효율적으로 계산하는 것입니다. 이산 시간 SSM의 임펄스 응답(Impulse Response)인 \bar{K}는 다음과 같은 수열로 정의됩니다.
\bar{K} = \left( \bar{C}\bar{B}, \bar{C}\bar{A}\bar{B}, \bar{C}\bar{A}^2\bar{B}, \dots, \bar{C}\bar{A}^{L-1}\bar{B} \right)

4.1 생성 함수(Generating Function) 접근법

이 수열을 직접 계산(Naive computation)하는 대신, S4는 생성 함수 접근법을 취합니다. 수열 \bar{K}z-변환인 생성 함수 \mathcal{K}(z)는 다음과 같이 표현됩니다.
\mathcal{K}(z) = \sum_{k=0}^{\infty} (\bar{C}\bar{A}^k\bar{B}) z^k = \bar{C} (I - \bar{A}z)^{-1} \bar{B}
여기서 z는 복소평면 상의 변수입니다. 만약 우리가 시퀀스 길이 L에 해당하는 단위근(Roots of Unity) \Omega_L = \{ e^{-2\pi i \frac{k}{L}} : k = 0, \dots, L-1 \} 위에서 \mathcal{K}(z)의 값을 빠르게 평가할 수 있다면, **역 고속 푸리에 변환(Inverse FFT)**을 통해 시간 영역의 커널 \bar{K}O(L \log L) 시간 안에 복원할 수 있습니다.

따라서 문제는 **“주어진 단위근 z에 대해 (I - \bar{A}z)^{-1}를 어떻게 빠르게 계산할 것인가?”**로 축소됩니다.

4.2 레졸번트(Resolvent)와 NPLR 구조

여기에 쌍선형 변환(Bilinear Transformation) 식을 대입하여 정리하면, 우리가 구해야 할 핵심 항은 다음과 같은 형태의 **레졸번트(Resolvent)**가 됩니다.
R(z) = \left( \frac{2}{\Delta} \frac{1-z}{1+z} I - A \right)^{-1}
S4의 핵심 아이디어는 행렬 ANPLR(Normal Plus Low Rank) 구조, 즉 A = \Lambda - P Q^* 형태를 가진다고 가정하는 것입니다(유니타리 행렬 V는 생략하거나 사전 처리되었다고 가정). 이를 대입하면 위 식은 다음과 같이 변형됩니다.
R(z) = (D(z) + P Q^*)^{-1}
여기서 D(z) = \frac{2}{\Delta} \frac{1-z}{1+z} I - \Lambda는 **대각 행렬(Diagonal Matrix)**입니다. 대각 행렬의 역행렬은 각 대각 성분의 역수만 취하면 되므로 O(N) 시간에 매우 쉽게 구할 수 있습니다.

4.3 우드버리 행렬 항등식(Woodbury Matrix Identity)의 적용

이제 우드버리 행렬 항등식을 사용할 차례입니다. 이 항등식은 (A + UCV)^{-1} 형태의 역행렬을 원래 행렬의 역행렬 A^{-1}과 낮은 차원의 수정항 조합으로 계산하게 해줍니다.
(D + PQ^*)^{-1} = D^{-1} - D^{-1} P (I + Q^* D^{-1} P)^{-1} Q^* D^{-1}
이 식의 계산 복잡도를 분석해보면 다음과 같습니다.

  • D^{-1}: 대각 행렬의 역행렬이므로 O(N)입니다.
  • Q^\* D^{-1} P: P, QN \times r 행렬이며, D^{-1}는 대각 행렬입니다. 이들의 곱은 r \times r 크기의 매우 작은 행렬이 됩니다 (일반적으로 r=1).
  • (I + \dots)^{-1}: r \times r 행렬의 역행렬 계산이므로 O(r^3) 또는 O(1)에 불과합니다.
  • 기타 행렬 곱셈: 모두 벡터 연산 또는 대각 행렬 연산으로 환원됩니다.

결과적으로, 우드버리 항등식은 N \times N 크기의 복잡한 행렬 역변환 문제를 대각 행렬 연산과 r \times r 소형 행렬 연산으로 분해함으로써 전체 복잡도를 **O(N)**으로 낮춥니다. 이것이 바로 S4가 거대한 상태 차원을 가지면서도 고속으로 연산할 수 있는 수학적 핵심입니다.

5. 코시 커널(Cauchy Kernel)로의 환원과 알고리즘 구현

우드버리 항등식을 적용한 후, 최종적으로 계산해야 하는 값은 입력 벡터(혹은 행렬 \bar{B})와 레졸번트의 곱이다. S4 논문에서는 이 연산이 **코시 행렬(Cauchy Matrix)**과 벡터의 곱셈 형태로 귀결됨을 보였다.

코시 행렬 C(\omega, \lambda)는 그 원소가 C_{ij} = \frac{1}{\omega_i - \lambda_j}로 정의되는 행렬이다. 여기서 \omega_i는 평가점(단위근 z)들이고, \lambda_j는 대각 행렬 \Lambda의 고유값들이다. 일반적인 행렬-벡터 곱셈은 O(N^2) 시간이 걸리지만, 코시 행렬은 그 특수한 구조 덕분에 더 빠른 알고리즘을 적용할 수 있다.

S4 구현에서는 시퀀스 길이 L개의 모든 위치(주파수)에서 동시에 값을 계산해야 한다. 이는 “빠른 다중극자 방법(Fast Multipole Method, FMM)“이나 “유리 함수 보간법(Rational Interpolation)“과 밀접하게 연관된 문제다. S4 연구진은 이를 해결하기 위해 FFT를 활용한 분할 정복 알고리즘을 제안하거나, 또는 단순히 GPU 상에서 최적화된 코시 커널 연산을 구현하여 사용했다.

이 과정을 정리하면 S4의 커널 생성 알고리즘은 다음과 같은 4단계 파이프라인으로 요약된다.

단계연산 내용복잡도비고
1. 파라미터 준비HiPPO 기반 A를 NPLR (\Lambda, P, Q, V)로 분해O(N)초기화 시 1회 수행
2. 주파수 영역 평가우드버리 항등식 및 코시 커널을 이용하여 L개의 단위근 z에 대해 생성 함수 \mathcal{K}(z) 계산O(N + L)핵심 가속 구간
3. 시간 영역 복원계산된 스펙트럼 값들에 대해 역 고속 푸리에 변환(iFFT) 수행O(L \log L)컨볼루션 커널 \bar{K} 획득
4. 컨볼루션 적용입력 시퀀스 u와 커널 \bar{K}의 FFT 기반 컨볼루션 (u * \bar{K})O(L \log L)실제 데이터 처리

이 알고리즘의 총 복잡도는 \tilde{O}(N + L)로, N에 대해 선형(Linear), L에 대해 로그-선형(Log-Linear)적이다. 이는 기존의 O(N^2 L)이나 O(L^2) 방식에 비해 압도적인 효율성을 자랑한다. 특히 N을 64, 128, 256 등으로 증가시켜도 학습 속도 저하가 거의 발생하지 않으므로, 모델은 매우 풍부한 상태 표현력을 가질 수 있게 된다.

6. S4D(Diagonal SSM)로의 진화와 단순화

S4의 NPLR 파라미터화와 우드버리 항등식은 수학적으로 완벽했지만, 실제 구현 난이도가 높고 수치적 안정성을 확보하기 위해 정교한 엔지니어링이 필요했다. 이에 후속 연구들은 “과연 NPLR 구조 전체가 꼭 필요한가?“라는 질문을 던졌다.

그 결과 등장한 것이 **S4D(Structured State Space Sequence Model with Diagonal State Matrices)**이다. S4D는 복잡한 P, Q 저순위 항을 제거하고, 순수하게 복소수 대각 행렬(Complex Diagonal Matrix) \Lambda만으로도 HiPPO 행렬의 성능을 근사할 수 있음을 실험적으로 증명했다.

S4D의 핵심 아이디어는 HiPPO 행렬을 대각화했을 때 얻어지는 고유값들의 분포를 분석하고, 이를 바탕으로 대각 행렬 \Lambda를 초기화하는 것이다. 비록 HiPPO 행렬 자체는 대각화가 불안정하지만, 그 고유값들이 갖는 기하학적 의미(복소평면의 좌반구에 위치하여 안정성을 보장하고 다양한 주파수 대역을 커버함)를 차용하여 대각 행렬을 직접 학습시키는 것이다.

S4D는 우드버리 항등식을 사용할 필요 없이 단순히 \bar{K} = \bar{C}(I - \bar{\Lambda}z)^{-1}\bar{B}를 계산하면 되므로 구현이 훨씬 간단하다. 이는 반데르몽드(Vandermonde) 행렬 곱셈과 유사한 형태가 되어 더욱 빠른 연산이 가능하다. 그러나 S4D의 성공은 S4가 “구조적 행렬을 이용한 대각화“라는 개념적 토대를 닦아놓았기에 가능했다. S4가 없었다면 연구자들은 HiPPO 행렬을 대각화하려는 시도조차 하지 않았을 것이기 때문이다.

7. 학습과 추론의 이원성(Duality): CNN처럼 학습하고 RNN처럼 추론한다

S4 대각화 기술의 궁극적인 가치는 학습과 추론 단계에서 서로 다른 뷰(View)를 취할 수 있다는 **이원성(Duality)**에 있다.

  1. 학습(Training) - 컨볼루션 뷰: 앞서 설명한 대로 전체 시퀀스를 한 번에 보는 병렬 처리가 가능하다. GPU의 수천 개 코어를 동시에 활용하여 긴 문맥을 순식간에 학습한다. 이는 트랜스포머의 병렬 학습 능력과 동등하다.
  2. 추론(Inference) - 순환 뷰: 학습이 완료된 후, 모델은 명시적인 컨볼루션 커널을 사용할 필요가 없다. 대신 이산화된 상태 행렬 \bar{A}, \bar{B}, \bar{C}를 이용하여 x_k = \bar{A}x_{k-1} + \bar{B}u_k 형태의 단순 행렬-벡터 곱셈을 수행한다.

이 이원성은 생성형 AI(Generative AI) 분야에서 결정적인 장점이 된다. 트랜스포머 기반의 LLM은 추론 시 KV 캐시(Key-Value Cache)를 유지해야 하며, 시퀀스 길이가 길어질수록 캐시 메모리 사용량과 연산량이 증가한다. 반면, S4 기반 모델(그리고 이를 계승한 Mamba)은 시퀀스 길이에 상관없이 항상 고정된 크기의 상태 x_k(N차원 벡터)만 유지하면 된다.

즉, S4는 O(1)의 상수 시간 추론(Constant-time Inference)과 상수 메모리(Constant Memory)를 보장한다. 이는 100만 토큰 길이의 문맥을 처리할 때도 메모리가 폭발하지 않고 일정한 속도로 텍스트를 생성할 수 있음을 의미하며, 엣지 디바이스나 실시간 애플리케이션에서의 활용 가능성을 극대화한다.

8. 소결: Mamba를 향한 디딤돌

요약하자면, “2.3 S4 대각화와 컨볼루션 기반 학습 가속“은 상태 공간 모델을 단순한 이론적 장난감에서 강력한 딥러닝 백본으로 격상시킨 기념비적인 기술이다.

  • HiPPO 행렬의 수치적 불안정성NPLR 분해를 통해 극복했다.
  • 행렬 역변환의 계산 비용우드버리 항등식코시 커널을 통해 획기적으로 낮췄다.
  • 학습 속도의 병목FFT 기반 컨볼루션으로 해소하여 트랜스포머급의 병렬성을 확보했다.
  • 그러면서도 추론 효율성은 RNN의 상수 시간 복잡도를 유지했다.

이러한 S4의 수학적 구조는 이후 등장할 Mamba 아키텍처의 기술적 모태가 된다. Mamba는 S4의 선형 시간 불변(LTI) 제약을 넘어 입력에 따라 파라미터가 변하는 선택적(Selective) 메커니즘을 도입하지만, 그 기저에 깔린 하드웨어 최적화 철학—메모리 계층 구조를 고려한 효율적 상태 관리와 스캔 알고리즘—은 S4의 연구 성과에 뿌리를 두고 있다. S4의 대각화 기법을 완벽히 이해하는 것은, 포스트 트랜스포머 시대의 새로운 패러다임인 SSM의 본질을 꿰뚫는 가장 확실한 길이다.